# -*- coding: utf-8 -*-
"""
Created on Mon Apr 12 08:35:40 2021

@author: Lenovo-pc
"""

import pandas as pd
import numpy as np
from sklearn import preprocessing


# numeric_path ='../data_gene/mdp_numeric.txt'
# traits_path = '../data_gene/mdp_traits.txt'

class get_mdp_data:
    def __init__(self,numeric_path,traits_path):
        self.numeric_path = numeric_path
        self.traits_path = traits_path
        
        self.numeric = pd.read_table(numeric_path) #原始numeric 带有taxa
        self.traits = pd.read_table(traits_path) #原始traits 带有Taxa
        #print(self.numeric.columns)
        self.id = self.traits['Taxa'] # 获得Taxa
        
        if('Taxa' in self.traits.columns):
            self.traits = self.traits.rename(columns={'Taxa':'taxa'})
            
        self.data = pd.merge(self.traits,self.numeric,on='taxa')
        self.data = self.data.dropna()
        self.numeric_cols = self.numeric.columns.tolist()
        self.numeric_cols.remove('taxa')
        
        self.traits_cols = self.traits.columns.tolist()
        
        
        self.X = self.data[self.numeric_cols]
        self.data_index = self.X.index.tolist()
        self.split()
        
        

    def split(self):
        # train_data= self.X.sample(frac=0.8,random_state=0)
        # test_data = self.X.drop(train_data.index)
        # self.train_index = train_data.index
        # self.test_index = test_data.index
    
        # self.train_labels = train_data.pop(self.traits_flag)
        # self.test_labels = test_data.pop(self.traits_flag)
        train_index = self.data_index 
        test_index=[]
       
        index_len = len(self.data_index)
        test_size = int(0.2*index_len)
        i=0
        while(i<test_size):
            randomIndex=int(np.random.uniform(0,index_len))
            if randomIndex in test_index:
                continue
            elif randomIndex not in train_index:
                continue
            else:
                test_index.append(randomIndex) 
                train_index.remove(randomIndex)
                i=i+1
                
        self.train_index = train_index
        self.test_index = test_index
            
        
    # def get_X(self):
    #     self.X = self.data[self.numeric_cols]
        
        
    def set_traits_flag(self,traits_flag):
        assert traits_flag in ['EarHT','dpoll','EarDia']  ,\
        "Error Input:common.get_mdp_data"
        self.traits_flag = traits_flag
        # numeric_cols_tmp = self.numeric_cols.copy()
        # numeric_cols_tmp.append(self.traits_flag)
        # self.df = self.data[numeric_cols_tmp]
        self.Y = self.data[self.traits_flag]
    
    def get_split_x(self):
        train_data = self.X.loc[self.train_index]
        test_data = self.X.loc[self.test_index]    
        return train_data.values,test_data.values
        
    def get_one_hot_x(self):
        train_data = self.X.loc[self.train_index]
        test_data = self.X.loc[self.test_index]
    
        categorieslist = [[0,1,2] for i in range(len(self.numeric_cols))]
        
        enc = preprocessing.OneHotEncoder(categories=categorieslist)
        enc.fit(train_data)
        train_data = enc.transform(train_data).toarray()
        test_data = enc.transform(test_data).toarray()
        return train_data,test_data
    
    def get_split_y(self):
        train_labels = self.Y.loc[self.train_index]
        test_labels = self.Y.loc[self.test_index]
        return train_labels.values,test_labels.values


if __name__ =="__main__":
    numeric_path ='../data/mdp_numeric.txt'
    traits_path = '../data/mdp_traits.txt'
    
    mdp_data = get_mdp_data(numeric_path,traits_path)
    
    train_data,test_data = mdp_data.get_one_hot_x()
    
    # traits_flag='EarHT'
    # mdp_data.set_traits_flag(traits_flag)
    # train_labels,test_labels  = mdp_data.get_split_y()
    
    
    # train_data,train_labels,test_data,test_labels = mdp_data.one_hot()
        
